import random

import numpy as np
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader
from P03_NER.LSTM_CRF.utils.common import *

datas, word2index = build_data()

class NerDataset(Dataset):
    def __init__(self, datas):
        super().__init__()
        self.datas = datas

    def __len__(self):
        return len(self.datas)

    def __getitem__(self, item):
        x = self.datas[item][0]
        y = self.datas[item][1]
        return x, y

def collate_fn(batch):
    # batch: [([],[]),([],[])……]
    x_train = []
    for data in batch:
        # 将文字转为id
        id_list = [word2index.get(word, 1) for word in data[0]]
        # 转成tensor张量
        x_train.append(torch.tensor(id_list))

    y_train = []
    for data in batch:
        # 将文字转为id
        id_list = [config.tag2id.get(word, 11) for word in data[1]]
        # 转成tensor张量
        y_train.append(torch.tensor(id_list))

    # 统一样本长度
    input_ids_padded = pad_sequence(x_train, batch_first=True, padding_value=0)
    labels_padded = pad_sequence(y_train, batch_first=True, padding_value=11)

    # 创建attention_mask
    attention_mask = (input_ids_padded != 0).long()
    return input_ids_padded, labels_padded, attention_mask

def get_data():
    # 随机打乱数据
    random.seed(24)
    random.shuffle(datas)
    # 划分训练集和验证集
    train_dataset = NerDataset(datas[:6300])
    train_dataloader = DataLoader(dataset=train_dataset,
                                  batch_size=config.batch_size,
                                  shuffle=True,
                                  drop_last=True,
                                  collate_fn=collate_fn)
    #  构造验证集加载器
    dev_dataset = NerDataset(datas[6300:])
    dev_dataloader = DataLoader(dataset=dev_dataset,
                                batch_size=config.batch_size,
                                shuffle=True,
                                drop_last=True,
                                collate_fn=collate_fn)

    return train_dataloader, dev_dataloader

if __name__ == '__main__':
    train_dataloader, dev_dataloader = get_data()
    for input_ids_padded, labels_padded, attention_mask in train_dataloader:
        print(f'input_ids_padded-->{input_ids_padded.shape}')
        print(f'labels_padded-->{labels_padded.shape}')
        print(f'attention_mask-->{attention_mask.shape}')
